"""
用于从gazebo生成的仿真图像中获取GT信息，边缘点，无拟合，无内点
"""

import numpy as np
import cv2
import os
import pickle
from scipy.spatial.transform import Rotation
from utils import SphereCoordinate
from typing import Iterable, TextIO
import tqdm


class ContourGT:
    def __init__(
        self,
        label_dir,
        origin_point: tuple,
        K: tuple,
        size: tuple,
        *,
        cache_dir=None,
        modifier,
        lat: tuple,
        lon: tuple,
        pixel_range: tuple,
        **kwargs,
    ):
        self.transform = SphereCoordinate(modifier, *lat, *lon, *pixel_range)
        if cache_dir is not None:
            if not os.path.exists(cache_dir):
                os.makedirs(cache_dir)
            if os.path.exists(os.path.join(cache_dir, "gt_cache.pkl")):
                with open(os.path.join(cache_dir, "gt_cache.pkl"), "rb") as f:
                    args_dict = pickle.load(f)
                    for k, v in args_dict.items():
                        setattr(self, k, v)
            else:
                dem, data, idx = self.preload(label_dir, **kwargs["preload"])
                self.origin_point = np.array(origin_point)
                self.K = np.array(K).reshape(3, 4)
                self.size = size
                self.points_, self.idx = self.findCraters(
                    dem, data, idx, **kwargs["findCraters"]
                )
                self.map_idx = idx
                with open(os.path.join(cache_dir, "gt_cache.pkl"), "wb") as f:
                    pickle.dump(self.__dict__, f)
        else:
            dem, data, idx = self.preload(label_dir, **kwargs["preload"])
            self.origin_point = np.array(origin_point)
            self.K = np.array(K).reshape(3, 4)
            self.size = size
            self.points_, self.idx = self.findCraters(
                dem, data, idx, **kwargs["findCraters"]
            )
            self.map_idx = idx

    @property
    def points(self):
        return self.points_

    def findCraters(self, dem, data, idx, *args, **kwargs):
        """
        从DEM中找到所有的陨石坑的边缘三维点，并将其坐标转为直角坐标
        """
        raise NotImplementedError

    def preload(
        self, label_dir, *args, **kwargs
    ) -> tuple[np.ndarray, np.ndarray, np.ndarray]:
        """从数据集中加载带标注的DEM数据，返回值依次是DEM数据，标注数据，标注数据的索引"""
        raise NotImplementedError

    def __call__(self, t, pose, ori, output_dir, img=None) -> bool:
        """
        输入的是相机在世界坐标系下的位置和姿态，需要输出的是当前位姿条件下陨石坑在图像坐标系下各点位置
        Arguments:
            t (str|float) : 时间戳
            pose (list) : Gazebo仿真输出的位置
            ori (list) : Gazebo仿真输出的欧拉角姿态，按照Z-Y-X顺序，pitch-yaw-roll顺序排列结果
            img (np.ndarray) : 用于显示的图像, optional
        """
        image_dir = os.path.join(output_dir, "gt_images")
        points_id_dir = os.path.join(output_dir, "points_id")
        if not os.path.exists(image_dir):
            os.makedirs(image_dir)
        if not os.path.exists(points_id_dir):
            os.makedirs(points_id_dir)
        if img is None:
            img = np.zeros(self.size, dtype=np.uint8)
        p = self.world2camera(pose, ori)
        # p, p_idx = self.clip(p)
        img, p, o_p, p_idx = self.fill_points(np.round(p).astype(np.int64), img)
        if p is None or p.size < 100:
            # 数量太少，标注图像没有意义，应当舍去
            return False
        cv2.imwrite(f"{image_dir}/{t}.png", img)
        with open(f"{points_id_dir}/{t}.txt", "w") as f:
            f.write(f"x/pix,y/pix,crater_id\n")
            for (x, y), id in zip(p.T, p_idx):
                f.write(f"{x},{y},{self.map_idx[id]}\n")
        return True

    def world2camera(self, pose, ori):
        # R = self.Rotate(ori)
        R = Rotation.from_quat(ori).as_matrix().T
        # R_ 调整相机坐标系和图像坐标系的坐标转换，即y轴x轴互换
        R_ = np.array([[0, 1, 0, 0], [1, 0, 0, 0], [0, 0, 1, 0], [0, 0, 0, 1]])
        T = -R @ np.array(pose)
        H = np.hstack((R, T[:, None]))
        H = np.vstack((H, np.array([0, 0, 0, 1])))
        homo_points = self.K @ R_ @ H @ self.points
        # 齐次坐标转绝对坐标
        p = homo_points[:2] / homo_points[-1]
        return p

    def fill_points(self, points, image):
        # clip
        x = points[0]
        y = points[1]
        ind = (x < self.size[0]) & (x >= 0) & (y < self.size[1]) & (y >= 0)
        idx = self.idx[ind]
        points = points[:, ind]
        origin_points = self.points[:, ind]
        points_clip = []
        origin_points_clip = []
        idx_clip = []
        if idx.size == 0:
            return None, None, None, None
        for it, i in enumerate(np.unique(idx)):
            points_i = points[:, idx == i]
            if self.diameter(points_i) > 0.5 * image.shape[0]:
                continue
            image[points_i[1], points_i[0], ...] = 255
            points_clip.append(points_i)
            origin_points_clip.append(origin_points)
            idx_clip.append(idx[idx == i])
        if it < 4:
            return None, None, None, None
        return (
            image,
            np.concatenate(points_clip, axis=1),
            np.concatenate(origin_points_clip, axis=1),
            np.concatenate(idx_clip),
        )

    def diameter(self, points):
        height = points[0].max() - points[0].min()
        width = points[1].max() - points[1].min()
        return max(height, width)


class SegmentGT:
    def __init__(self, size, label_dir):
        with open(os.path.join(label_dir, "gt_labels.txt"), "r") as f:
            self.ind = map(lambda x: x.strip(), f.readlines())
        self.label_dir = label_dir
        # self.root_name = os.path.basename(label_dir)
        self.size = size
        self.output_dir = os.path.join(label_dir, "gt_segment")
        if not os.path.exists(self.output_dir):
            os.makedirs(self.output_dir)

    def __enter__(self):
        self.writer = open(os.path.join(self.label_dir, f"gt_segment.txt"), "w")
        return self

    def __exit__(self, exc_type, exc_val, exc_tb):
        self.writer.close()
        # raise exc_type

    def __call__(self, use_bg: bool):
        for ind in tqdm.tqdm(self.ind, desc="Generating segment images"):
            # 对于每一张图像，读取标注点信息
            if use_bg:
                img = cv2.imread(os.path.join(self.label_dir, "images", f"{ind}.png"))
            else:
                img = np.zeros(self.size, dtype=np.uint8)
            with open(os.path.join(self.label_dir, "points_id", f"{ind}.txt")) as f:
                points_iter = self._read_label(f)
                img, th_num = self._fill_hole(points_iter, img)
            # 保存图像
            if th_num:
                print(ind, file=self.writer)
                cv2.imwrite(os.path.join(self.output_dir, f"{ind}.png"), img)

    def _read_label(self, ind: TextIO):
        """
        流式读取标注点的信息
        """
        _ = ind.readline()
        for i, line in enumerate(ind):
            x, y, idx = line.split(",")
            if i == 0:
                points = [(int(x), int(y))]
                idx_last = idx
            elif idx == idx_last:
                points.append((int(x), int(y)))
            else:
                yield points
                points = [(int(x), int(y))]
                idx_last = idx
        # 输出最后一个点列表
        yield points

    def _fill_hole(self, points_iter, mask):
        """
        用洪水淹没法填充孔洞，TODO 在边界上有部分孔洞无法填充
        """
        output = mask.copy()
        for points in map(np.array, points_iter):
            if points.shape[0] < 100:
                continue
            temp = np.zeros_like(mask, dtype=np.uint8)
            temp[points[:, 1], points[:, 0]] = 255
            temp = cv2.morphologyEx(
                temp,
                cv2.MORPH_CLOSE,
                cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (5, 5)),
            )
            cc_num, temp_mask = cv2.connectedComponents(temp)
            if cc_num > 2:
                continue
            cv2.floodFill(
                temp,
                np.pad(temp_mask.astype(np.uint8), 1, "constant", constant_values=1),
                (int(points[:, 0].mean()), int(points[:, 1].mean())),
                255,
            )
            output = cv2.bitwise_or(output, temp)
        if output.std() < 10 or (output != 0).sum() < 500:
            return output, False
        return output, True

    def _fit_ellipse(
        self, points_iter: Iterable, img: np.ndarray, tol=4, th_num=3
    ) -> np.ndarray:
        """
        拟合椭圆，TODO 但是在个别地方会出现误拟合的情况
        Arguments:
            points_iter (Iterable) : an iterator for points list with same id
            img (np.ndarray) : the image to fill ellipse
            tol (float) : the tolerance of the ratio of major axis and minor axis
            th_num (int) : the minimal number of points to fill ellipse
        """
        num_ellipse = 0
        for points in map(np.array, points_iter):
            if points.shape[0] > 10:  # 排除边缘点太少的陨石坑，避免误填充
                # 填充椭圆
                ellipse_box = cv2.fitEllipseAMS(points)
                # 长宽比例过大的椭圆不填充
                if (
                    ellipse_box[1][0] / ellipse_box[1][1] < 1 / tol
                    or ellipse_box[1][0] / ellipse_box[1][1] > tol
                ):
                    continue
                img = cv2.ellipse(img, ellipse_box, (255, 255, 255), -1)
                num_ellipse += 1

        return img, num_ellipse > th_num
